#include "kvalloc.h"
#include <cmath>

namespace kvalloc {

PMChunk::PMChunk(const std::string & fname, uint64_t size) : size_(size) {
    size_t mem_len;
    int is_pmem;
    if(file_exist(fname.c_str())) {
        // open the existing pool file and mmap it into virtual memory
        buf_ = pmem_map_file(fname.c_str(), 0, 0, 0666, &mem_len, &is_pmem); 
        if(!buf_) {
            perror("pmem_map_file error");
            exit(1);
        }
    } else {
        // create a new pool file and mmap it into virtual memory 
        buf_ = pmem_map_file(fname.c_str(), size_, PMEM_FILE_CREATE, 0666, &mem_len, &is_pmem);
        if(!buf_) {
            perror("pmem_map_file error");
            exit(1);
        }
        // zero the chunk
        memset(buf_, 0, size_);
    }

    if(mem_len < size_) {
        printf("The original file is too small");
        exit(1);
    }

    // whether the pool file locates on PMEM or normal file system? 
    if(!is_pmem) {
        printf("This file is not on pmem");
        exit(1);
    }
}

PMChunk::~PMChunk() {
    pmem_unmap(buf_, size_);
}

KVAlloc::KVAlloc(const std::string & dbfname, const std::string & metafname, uint64_t max_snum)
    : max_snum_(max_snum), 
    file_(dbfname.c_str()), 
    chunk_(metafname, sizeof(NVMetaWrapper) + max_snum * sizeof(SlabMeta)) 
{   
    // Persistent states of KVAlloc is allocated from NVM for persistence
    nvmeta_ = (NVMetaWrapper *)chunk_.GetBuf();
}

FileAddr KVAlloc::Alloc(size_t size) {
    FileAddr res;
    auto iter = SIZECLASS.lower_bound(size);
    int size_class = iter->first;
    int sc_id = iter->second; // size class id

    SlabMeta * front;
    // 锁住对应链表
    nvmeta_->alloc_sm[sc_id].mu_.Lock();
    alloc_retry_locked:
        nvmeta_->alloc_sm[sc_id].count_ += 1;
        // 偶发地将free_sm中的slab对象转移到alloc_sm，减少alloc和free的冲突
        // 每从当前链表分配满100次时，抓取一次对应的free_slabs
        static constexpr uint32_t SAMPLE   = 100; 
        static constexpr uint32_t MIN_GRAB = 1; 
        // count_ 在alloc_sm表示分配次数，在free_sm中表示链表长度
        if(nvmeta_->alloc_sm[sc_id].count_ % SAMPLE == 0 
                && nvmeta_->free_sm[sc_id].count_ >= MIN_GRAB) {
            nvmeta_->free_sm[sc_id].mu_.Lock();
                nvmeta_->alloc_sm[sc_id].Append(nvmeta_->all_sm, nvmeta_->free_sm[sc_id]); // 将free_sm链表加入到alloc_sm头部
                nvmeta_->free_sm[sc_id].Clear(); // 清空free_sm
            nvmeta_->free_sm[sc_id].mu_.UnLock();
        }

        // 获取分配链表第一个slab
        front = nvmeta_->alloc_sm[sc_id].Front(nvmeta_->all_sm);
        // 第一个页用来作为哨兵
        if(front->slabid_ == 0) { // 如果当前链表为空，从文件中分配新页当成slab
            size_t cur_pid = file_.AllocPages(ALLOC_BATCH);
            assert(cur_pid + ALLOC_BATCH <= max_snum_);
            for(int i = 0; i < ALLOC_BATCH; i++) {
                SlabMeta * sm = &(nvmeta_->all_sm[cur_pid + i]);
                sm->size_class_ = size_class;
                sm->slabid_ = cur_pid + i; // 首次分配，初始化slabid
                sm->state_ = NONFULL;
                nvmeta_->alloc_sm[sc_id].PushBack(nvmeta_->all_sm, cur_pid + i);
            }
            front = nvmeta_->alloc_sm[sc_id].Front(nvmeta_->all_sm);
        }
    // 释放当前链表
    nvmeta_->alloc_sm[sc_id].mu_.UnLock();

    bool success = front->Alloc(res); // alloc a slot
    if(!success) { // 分配失败，说明front这个slab已经满了
        nvmeta_->alloc_sm[sc_id].mu_.Lock();
        if(nvmeta_->alloc_sm[sc_id].Front(nvmeta_->all_sm) == front 
                && front->state_ == NONFULL) {
            // 更新slab状态
            front->state_ = FULL;
            // 从当前链表中删除Front
            nvmeta_->alloc_sm[sc_id].PopFront(nvmeta_->all_sm);
        }
        goto alloc_retry_locked;
    } else {
        return res; // 分配成功
    }
}

void KVAlloc::Free(FileAddr addr) {
    size_t slabid = addr.slabid;
    SlabMeta * sm = &(nvmeta_->all_sm[slabid]);
    sm->Free(addr); // free a slot
    if(sm->state_ == FULL && sm->TryLock()) { // 如果当前slab为满状态，则需要有一个线程负责将其加入到free_sm
        // 更新slab状态
        sm->state_ = NONFULL;  
        // 加入对应的free_sm链表
        int sc_id = SIZECLASS.at(sm->size_class_); // size class id
        nvmeta_->free_sm[sc_id].mu_.Lock();
            nvmeta_->free_sm[sc_id].PushBack(nvmeta_->all_sm, slabid);
            nvmeta_->free_sm[sc_id].count_ += 1;
        nvmeta_->free_sm[sc_id].mu_.UnLock();
        // 赢家处理完毕，释放锁
        sm->UnLock();
    }
}

void KVAlloc::Read(char * buf, FileAddr addr) {
    size_t slabid = addr.slabid;
    SlabMeta * sm = &(nvmeta_->all_sm[slabid]);
    uint64_t size = sm->size_class_;
    file_.Read(addr.FileOff(), buf, size);
}

size_t KVAlloc::GetSize(FileAddr addr) {
    size_t slabid = addr.slabid;
    SlabMeta * sm = &(nvmeta_->all_sm[slabid]);
    return sm->size_class_;
}

} // kvalloc 